
import whisper
import opencc
import numpy as np
from flask import Flask, request, jsonify
import os
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = './uploads'
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 初始化模型和转换器（全局加载避免重复初始化）
model = whisper.load_model("large-v3")
cc = opencc.OpenCC('t2s.json')

def process_audio(audio_path):
    result = model.transcribe(audio_path, word_timestamps=True)
    result['text'] = cc.convert(result['text'])
    
    segments = result['segments']
    for segment in segments:
        segment.update({
            'start': f"{np.round(segment['start'], 2)}",
            'end': f"{np.round(segment['end'], 2)}",
            'text': cc.convert(segment['text'])
        })
        for key in ['tokens','words','temperature','avg_logprob','compression_ratio','no_speech_prob']:
            segment.pop(key, None)
    return result

@app.route('/transcribe', methods=['POST'])
def transcribe_api():
    if 'file' not in request.files:
        return jsonify({'code': 400, 'error': 'No file uploaded'})
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'code': 400, 'error': 'Empty filename'})
    
    try:
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(save_path)
        
        result = process_audio(save_path)
        result['code'] = 0
        os.remove(save_path)  # 处理完成后删除临时文件
        
        return jsonify(result)
    except Exception as e:
        return jsonify({'code': 500, 'error': str(e)})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
